Fast Sampling of Diffusion Models via Operator Learning

Diffusion models have found widespread adoption in various areas. However, sampling from them is slow because it involves emulating a reverse process with hundreds-to-thousands of network evaluations. Inspired by the success of neural operators in accelerating differential equations solving, we approach this problem by solving the underlying neural differential equation from an operator learning perspective. We examine probability flow ODE trajectories in diffusion models and observe a compact energy spectrum that can be learned efficiently in Fourier space. With this insight, we propose diffusion Fourier neural operator (DFNO) with temporal convolution in Fourier space to parameterize the operator that maps initial condition to the solution trajectory, which is a continuous function in time. DFNO can be applied to any diffusion model and generate high-quality samples in one model forward call. Our method achieves the state-of-the-art FID of 4.72 on CIFAR-10 using only one model evaluation

Improving Generative Model-based Unfolding with Schr\"{o}dinger Bridges

Machine learning-based unfolding has enabled unbinned and high-dimensional differential cross section measurements. Two main approaches have emerged in this research area: one based on discriminative models and one based on generative models. The main advantage of discriminative models is that they learn a small correction to a starting simulation while generative models scale better to regions of phase space with little data. We propose to use Schroedinger Bridges and diffusion models to create SBUnfold, an unfolding approach that combines the strengths of both discriminative and generative models. The key feature of SBUnfold is that its generative model maps one set of events into another without having to go through a known probability density as is the case for normalizing flows and standard diffusion models. We show that SBUnfold achieves excellent performance compared to state of the art methods on a synthetic Z+jets dataset.Comment: 9 pages, 5 figure

Dynamic-Backbone Protein-Ligand Structure Prediction with Multiscale Generative Diffusion Models

Molecular complexes formed by proteins and small-molecule ligands are ubiquitous, and predicting their 3D structures can facilitate both biological discoveries and the design of novel enzymes or drug molecules. Here we propose NeuralPLexer, a deep generative model framework to rapidly predict protein-ligand complex structures and their fluctuations using protein backbone template and molecular graph inputs. NeuralPLexer jointly samples protein and small-molecule 3D coordinates at an atomistic resolution through a generative model that incorporates biophysical constraints and inferred proximity information into a time-truncated diffusion process. The reverse-time generative diffusion process is learned by a novel stereochemistry-aware equivariant graph transformer that enables efficient, concurrent gradient field prediction for all heavy atoms in the protein-ligand complex. NeuralPLexer outperforms existing physics-based and learning-based methods on benchmarking problems including fixed-backbone blind protein-ligand docking and ligand-coupled binding site repacking. Moreover, we identify preliminary evidence that NeuralPLexer enriches bound-state-like protein structures when applied to systems where protein folding landscapes are significantly altered by the presence of ligands. Our results reveal that a data-driven approach can capture the structural cooperativity among protein and small-molecule entities, showing promise for the computational identification of novel drug targets and the end-to-end differentiable design of functional small-molecules and ligand-binding proteins

Retrieval-based Controllable Molecule Generation

Generating new molecules with specified chemical and biological properties via generative models has emerged as a promising direction for drug discovery. However, existing methods require extensive training/fine-tuning with a large dataset, often unavailable in real-world generation tasks. In this work, we propose a new retrieval-based framework for controllable molecule generation. We use a small set of exemplar molecules, i.e., those that (partially) satisfy the design criteria, to steer the pre-trained generative model towards synthesizing molecules that satisfy the given design criteria. We design a retrieval mechanism that retrieves and fuses the exemplar molecules with the input molecule, which is trained by a new self-supervised objective that predicts the nearest neighbor of the input molecule. We also propose an iterative refinement process to dynamically update the generated molecules and retrieval database for better generalization. Our approach is agnostic to the choice of generative models and requires no task-specific fine-tuning. On various tasks ranging from simple design criteria to a challenging real-world scenario for designing lead compounds that bind to the SARS-CoV-2 main protease, we demonstrate our approach extrapolates well beyond the retrieval database, and achieves better performance and wider applicability than previous methods.Comment: 29 page

Bongard-LOGO: A New Benchmark for Human-Level Concept Learning and Reasoning

Humans have an inherent ability to learn novel concepts from only a few samples and generalize these concepts to different situations. Even though today's machine learning models excel with a plethora of training data on standard recognition tasks, a considerable gap exists between machine-level pattern recognition and human-level concept learning. To narrow this gap, the Bongard Problems (BPs) were introduced as an inspirational challenge for visual cognition in intelligent systems. Albeit new advances in representation learning and learning to learn, BPs remain a daunting challenge for modern AI. Inspired by the original one hundred BPs, we propose a new benchmark Bongard-LOGO for human-level concept learning and reasoning. We develop a program-guided generation technique to produce a large set of human-interpretable visual cognition problems in action-oriented LOGO language. Our benchmark captures three core properties of human cognition: 1) context-dependent perception, in which the same object may have disparate interpretations given different contexts; 2) analogy-making perception, in which some meaningful concepts are traded off for other meaningful concepts; and 3) perception with a few samples but infinite vocabulary. In experiments, we show that the state-of-the-art deep learning methods perform substantially worse than human subjects, implying that they fail to capture core human cognition properties. Finally, we discuss research directions towards a general architecture for visual reasoning to tackle this benchmark

Bongard-LOGO: A New Benchmark for Human-Level Concept Learning and Reasoning

Humans have an inherent ability to learn novel concepts from only a few samples and generalize these concepts to different situations. Even though today's machine learning models excel with a plethora of training data on standard recognition tasks, a considerable gap exists between machine-level pattern recognition and human-level concept learning. To narrow this gap, the Bongard problems (BPs) were introduced as an inspirational challenge for visual cognition in intelligent systems. Despite new advances in representation learning and learning to learn, BPs remain a daunting challenge for modern AI. Inspired by the original one hundred BPs, we propose a new benchmark Bongard-LOGO for human-level concept learning and reasoning. We develop a program-guided generation technique to produce a large set of human-interpretable visual cognition problems in action-oriented LOGO language. Our benchmark captures three core properties of human cognition: 1) context-dependent perception, in which the same object may have disparate interpretations given different contexts; 2) analogy-making perception, in which some meaningful concepts are traded off for other meaningful concepts; and 3) perception with a few samples but infinite vocabulary. In experiments, we show that the state-of-the-art deep learning methods perform substantially worse than human subjects, implying that they fail to capture core human cognition properties. Finally, we discuss research directions towards a general architecture for visual reasoning to tackle this benchmark.Comment: 22 pages, NeurIPS 202

Mapping Region-Specific Longitudinal Cortical Surface Expansion from Birth to 2 Years of Age

The human cerebral cortex develops rapidly and dynamically in the first 2 years of life. It has been shown that cortical surface expansion from term infant to adult is highly nonuniform in a cross-sectional study. However, little is known about the longitudinal cortical surface expansion during early postnatal stages. In this article, we generate the first longitudinal surface-based atlases of human cortical structures at 0, 1, and 2 years of age from 73 healthy subjects. On the basis of the surface-based atlases, we study the longitudinal cortical surface expansion in the first 2 years of life and find that cortical surface expansion is age related and region specific. In the first year, cortical surface expands dramatically, with an average expansion of 1.80 times. In particular, regions of superior and medial temporal, superior parietal, medial orbitofrontal, lateral anterior prefrontal, occipital cortices, and postcentral gyrus expand relatively larger than other regions. In the second year, cortical surface still expands substantially, with an average expansion of 1.20 times. In particular, regions of superior and middle frontal, orbitofrontal, inferior temporal, inferior parietal, and superior parietal cortices expand relatively larger than other regions. These region-specific patterns of cortical surface expansion are related to cognitive and functional development at these stages

Measuring the dynamic longitudinal cortex development in infants by reconstruction of temporally consistent cortical surfaces

Quantitative measurement of the dynamic longitudinal cortex development during early postnatal stages is of great importance to understand the early cortical structural and functional development. Conventional methods usually reconstruct the cortical surfaces of longitudinal images from the same subject independently, which often generate longitudinally-inconsistent cortical surfaces and thus lead to inaccurate measurement of cortical changes, especially for vertex-wise mapping of cortical development. This paper aims to address this problem by presenting a method to reconstruct temporally-consistent cortical surfaces from longitudinal infant brain MR images, for accurate and consistent measurement of the dynamic cortex development in infants. Specifically, the longitudinal development of the inner cortical surface is first modeled by a deformable growth sheet with elasto-plasticity property to establish longitudinally smooth correspondences of the inner cortical surfaces. Then, the modeled longitudinal inner cortical surfaces are jointly deformed to locate both inner and outer cortical surfaces with a spatial-temporal deformable surface method. The method has been applied to 13 healthy infants, each with 6 serial MR scans acquired at 2 weeks, 3 months, 6 months, 9 months, 12 months and 18 months of age. Experimental results showed that our method with the incorporated longitudinal constraints can reconstruct the longitudinally-dynamic cortical surfaces from serial infant MR images more consistently and accurately than the previously published methods. By using our method, for the first time, we can characterize the vertex-wise longitudinal cortical thickness development trajectory at multiple time points in the first 18 months of life. Specifically, we found the highly age-related and regionally-heterogeneous developmental trajectories of the cortical thickness during this period, with the cortical thickness increased most from 3 to 6 months (16.2%) and least from 9 to 12 months (less than 0.1%). Specifically, the central sulcus only underwent significant increase of cortical thickness from 6 to 9 months and the occipital cortex underwent significant increase from 0 to 9 months, while the frontal, temporal and parietal cortices grew continuously in this first 18 months of life. The adult-like spatial patterns of cortical thickness were generally present at 18 months of age. These results provided detailed insights into the dynamic trajectory of the cortical thickness development in infants